import tensorflow as tf
from tensorflow.keras.preprocessing.image import save_img
from stylegan2_pytorch.model import Generator, Discriminator

# 加载预训练模型
generator = Generator(size=256, style_dim=512, n_mlp=8)
ckpt = tf.train.Checkpoint(generator)
ckpt.restore('path_to_your_checkpoint').expect_partial()

# 生成随机种子
seed = tf.random.normal([1, 512])

# 生成图片
img = generator(seed, training=False)

# 保存图片
save_img('generated_image.png', img[0])
